/**
 * 
 */
package coder;

import java.util.Arrays;

import utilities.ByteIntegerConverter;
import utilities.MinPQ;

/**
 * @author Marcin Wachulski
 *
 */
public final class HuffmanCoder implements Coder 
{
	private static final double CAPACITY_FACTOR 	= 	1.5;
	// alphabet size of extended ASCII
    private static final int 	R 					= 	256;
    // help constant to represent the largest one-bit ubyte
    private static final int 	MAX_NUMBIT 			= 	128;
	
    byte[] output;
    byte[] input;
    int outputIdx;
    int inputIdx;
	int nextBit;
    
	@Override
	public byte[] encode(byte[] input, int length) 
	{
		output = new byte[length + 4];
		outputIdx = 0;
		nextBit = MAX_NUMBIT;

        // tabulate frequency counts
        int[] freq = new int[R];
        for (int i = 0; i < length; i++)
            freq[((int)input[i] & 255)]++;

        // build Huffman trie
        Node root = buildTrie(freq);

        // build code table
        String[] st = new String[R];
        buildCode(st, root, "");

        // print number of symbols due to encode
        byte[] symbolsLength = ByteIntegerConverter.intToByteArray(length);
        output[outputIdx++] = symbolsLength[0];
        output[outputIdx++] = symbolsLength[1];
        output[outputIdx++] = symbolsLength[2];
        output[outputIdx++] = symbolsLength[3];
        
        // print trie for decoder
        writeTrie(root);

        // use Huffman code to encode input
        for (int i = 0; i < length; i++) {
            String code = st[((int) input[i]) & 255];
            for (int j = 0; j < code.length(); j++)
                if (code.charAt(j) == '0')
                    writeBit(false);
                else
                	writeBit(true);
        }
        if (nextBit < MAX_NUMBIT)
        	++outputIdx;
        
        // return decoded byte array
        return Arrays.copyOf(output, outputIdx);
	}

	@Override
	public byte[] decode(byte[] input) 
	{
		this.input = input;
		inputIdx = 0;
		output = new byte[input.length];
		outputIdx = 0;
		nextBit = MAX_NUMBIT;

        // number of symbols to decode
        int length = 0;
        length = ByteIntegerConverter.byteArrayToInt(Arrays.copyOfRange(input, 0, 4), 0);
        inputIdx += 4;
        
     // read in Huffman trie from input stream
        Node root = readTrie(); 
        
        // decode using the Huffman trie
        for (int i = 0; i < length; i++) {
            Node x = root;
            while (!x.isLeaf()) {
                boolean bit = readBit();
                if (bit) x = x.right;
                else     x = x.left;
            }
            output[outputIdx++] = x.symbol;
            if (output.length == outputIdx)
    			output = Arrays.copyOf(output, (int)((double)output.length * CAPACITY_FACTOR) + 1);
        }
		
		return Arrays.copyOf(output, outputIdx);
	}

    

    // Huffman trie node
    private static class Node implements Comparable<Node> {
        private final byte symbol;
        private final int freq;
        private final Node left, right;

        Node(byte symbol, int freq, Node left, Node right) {
            this.symbol	= symbol;
            this.freq  	= freq;
            this.left  	= left;
            this.right 	= right;
        }

        // is the node a leaf node?
        private boolean isLeaf() {
            assert (left == null && right == null) || (left != null && right != null);
            return (left == null && right == null);
        }

        // compare, based on frequency
        public int compareTo(Node that) {
            return this.freq - that.freq;
        }
    }

    // build the Huffman trie given frequencies
    private Node buildTrie(int[] freq) {
        // initialze priority queue with singleton trees
        MinPQ<Node> pq = new MinPQ<Node>();
        for (int i = 0; i < R; i++)
            if (freq[i] > 0)
                pq.insert(new Node((byte)i, freq[i], null, null));

        // merge two smallest trees
        while (pq.size() > 1) {
            Node left  = pq.delMin();
            Node right = pq.delMin();
            Node parent = new Node((byte)0, left.freq + right.freq, left, right);
            pq.insert(parent);
        }
        return pq.delMin();
    }


    // write bitstring-encoded trie to standard output
    private void writeTrie(Node x) {
        if (x.isLeaf()) {
            writeBit(true);
            writeByte(x.symbol);
            return;
        }
        writeBit(false);
        writeTrie(x.left);
        writeTrie(x.right);
    }

    // make a lookup table from symbols and their encodings
    private void buildCode(String[] st, Node x, String s) {
        if (!x.isLeaf()) {
            buildCode(st, x.left,  s + '0');
            buildCode(st, x.right, s + '1');
        }
        else {
            st[((int) x.symbol) & 255] = s;
        }
    }

    private Node readTrie() {
        boolean isLeaf = readBit();
        if (isLeaf) {
            return new Node(readByte(), -1, null, null);
        }
        else {
            return new Node((byte) 0, -1, readTrie(), readTrie());
        }
    }
    
    private void writeBit(boolean bit) {
    	if (bit)
    		output[outputIdx] |= nextBit;
    	if (nextBit == 1) {
			nextBit = MAX_NUMBIT;
			++outputIdx;
			if (output.length == outputIdx)
    			output = Arrays.copyOf(output, (int)((double)output.length * CAPACITY_FACTOR) + 1);
		}
		else
			nextBit >>>= 1;
    }
    
    private void writeByte(byte b) {
    	for (int i = 0; i < 8; ++i)
    		writeBit(((((int)b & 255) >> i) & 1) == 1);    		
    }
    
    private boolean readBit() {
    	boolean bit = (input[inputIdx] & nextBit) != 0;
    	if (nextBit == 1) {
    		nextBit = MAX_NUMBIT;
    		++inputIdx;
    	}
    	else
    		nextBit >>>= 1;
    	return bit;
    }
    
    private byte readByte() {
    	int b = 0;
    	for (int i = 0; i < 8; ++i)
    		b |= (readBit() ? 1 : 0) << i;
    	return (byte)b;
    }
}
